from .db_api import DbApi
import pymysql


class MySqlPlugin(DbApi):
    conn_type = 'mysql'
    _test_connection_sql = "select 1"

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.schema = kwargs.pop("schema", None)
        self.connection = kwargs.pop("connection", None)

    def get_conn(self):
        conn = self.connection or self.get_connection(getattr(self, self.conn_name_attr))
        conn_config = {
            'user': conn.login,
            'password': conn.password or '',
            'host': conn.host or 'localhost',
            'database': self.schema or conn.schema or '',
            'port': int(conn.port) if conn.port else 3306,
        }

        if conn.extra_dejson.get('charset', False):
            conn_config["charset"] = conn.extra_dejson["charset"]
            if conn_config["charset"].lower() in ('utf8', 'utf-8'):
                conn_config["use_unicode"] = True

        return pymysql.connect(**conn_config)

    def _get_cursor_by_type(self, conn, cursor=None):
        if cursor is None:
            return conn.cursor(pymysql.cursors.Cursor)
        elif cursor.lower() == "dictcursor":
            return conn.cursor(pymysql.cursors.DictCursor)
        else:
            return conn.cursor(pymysql.cursors.Cursor)
